import os
import json
import datetime

import numpy as np

import tqdm

from common.geometry.geom import Geom
from common.geoimage.gdal_sys import stack_ds, mosaic_ds
from common.geoimage.raster_dataset import (
    RasterDataset,
    SourceWindow,
    mosaic_raster_datasets,
)

from catalog.utils import interpret_protocol
from catalog.models.sentinel2_l2a import Sentinel2L2AScene


BASE_PATH = "/media/workspace/zhouhang/sample_data"


def test_open_single_band_image():
    dir_base = os.path.join(
        BASE_PATH, "S2B_MSIL2A_20190204T034949_N0211_R104_T48RVV_20190204T073411.SAFE"
    )
    fpath_input = os.path.join(
        dir_base,
        "GRANULE/L2A_T48RVV_A009997_20190204T034943/IMG_DATA/R10m/T48RVV_20190204T034949_B08_10m.jp2",
    )

    print("read a single band sample file")
    rds: RasterDataset = RasterDataset.from_file(fpath_input)
    print(rds.data.shape, rds.data.dtype)
    print(rds.meta)

    print("subset the image by window in image pixel space")
    subset_rds: RasterDataset = rds.subset(SourceWindow(100, 300, 500, 400))
    print(subset_rds.data.shape)
    print(subset_rds.meta)
    subset_rds.to_geotiff("test1_subset_by_window.tif")

    print("subset the image by window with a given aoi in wgs84")
    fpath_loc = os.path.join(
        BASE_PATH, "china_geojson_eng/xian4/guang3han4shi4.geojson"
    )
    info = json.loads(open(fpath_loc).read())
    aoi = Geom(info["features"][0]["geometry"])
    subset_rds1: RasterDataset = rds.subset(aoi=aoi)
    print(subset_rds1.data.shape)
    print(subset_rds1.meta)
    subset_rds1.to_jp2("test1_subset_by_aoi.jp2")


def test_open_multiple_band_image():
    fpath_input = os.path.join(
        BASE_PATH, "decloud_10m_51TXN_20200729_20200827_False.tif"
    )

    print("read them all")
    rds: RasterDataset = RasterDataset.from_file(fpath_input)
    print(rds.data.dtype, rds.data.shape)
    print(rds.meta)

    print("slice bands from the loaded raster dataset")
    rds_subset = rds.slice(bands=[1, 2], copy_ndarray=False)
    print(rds_subset.data.dtype, rds_subset.data.shape)
    print(rds_subset.meta)
    print(np.all(rds_subset.data[0] == rds.data[1]))
    print(np.all(rds_subset.data[1] == rds.data[2]))

    print("only read designated bands")
    rds1: RasterDataset = RasterDataset.from_file(fpath_input, bands=[0, 3])
    print(rds1.data.dtype, rds1.data.shape)
    print(rds1.meta)
    print(np.all(rds1.data[0] == rds.data[0]))
    print(np.all(rds1.data[1] == rds.data[3]))


def test_open_multiple_band_image_real_case():
    print(
        "read R and NIR bands from a multi-band image for a given aoi, then calculate NDVI and save as a cog"
    )

    fpath_input = os.path.join(
        BASE_PATH, "decloud_10m_51TXN_20200729_20200827_False.tif"
    )

    fpath_loc = os.path.join(
        BASE_PATH, "china_geojson_eng/shi4/qi2qi2ha1er3shi4.geojson"
    )
    info = json.loads(open(fpath_loc).read())
    aoi = Geom(info["features"][0]["geometry"])

    rds: RasterDataset = RasterDataset.from_file(fpath_input, bands=[0, 3], aoi=aoi)
    print(rds.meta)
    print(rds.data.shape, rds.data.dtype)
    rds.to_jp2("test3_partial_read.jp2")


def test_iter_from_file(simple=True):
    if simple:
        print("read from a per-tile image by block of 256X256")
        fpath_input = os.path.join(
            BASE_PATH, "decloud_10m_51TXN_20200729_20200827_False.tif"
        )
        block_size = 256
    else:
        print(
            "read from a big whole China NDVI image, and we have to do it by block of 4096X4096"
        )
        print("only traverse the first 50")
        fpath_input = os.path.join(BASE_PATH, "20200806_TO_20200905_all.tif")  # 72G
        block_size = 4096
    rds_traverser = RasterDataset.iter_from_file(
        fpath_input, block_x_size=block_size, block_y_size=block_size
    )
    print(rds_traverser.estimated_n_blocks)
    print(rds_traverser.n_blocks)
    for ind, (rds, block_info, _) in tqdm.tqdm(
        enumerate(rds_traverser.traverse()), total=rds_traverser.n_blocks
    ):
        if (not simple) and ind == 50:
            break


def test_warp_to_wgs84():
    dir_base = os.path.join(
        BASE_PATH, "S2B_MSIL2A_20190204T034949_N0211_R104_T48RVV_20190204T073411.SAFE"
    )
    fpath_input = os.path.join(
        dir_base,
        "GRANULE/L2A_T48RVV_A009997_20190204T034943/IMG_DATA/R10m/T48RVV_20190204T034949_B08_10m.jp2",
    )

    print("read a single band sample file")
    rds: RasterDataset = RasterDataset.from_file(fpath_input)
    print(rds.meta)

    rds_warped = rds.warp("EPSG:4326")
    print(rds_warped.meta)
    print(rds_warped.data)
    rds_warped.to_geotiff("./test5_b08_wgs84.tif")

    print("subset the image by window with a given aoi in wgs84")
    fpath_loc = os.path.join(
        BASE_PATH, "china_geojson_eng/xian4/guang3han4shi4.geojson"
    )
    info = json.loads(open(fpath_loc).read())
    aoi = Geom(info["features"][0]["geometry"])
    rds_warped1 = rds.warp(
        "EPSG:4326", cutline=aoi, crop_to_cutline=True, resample_alg="lanczos"
    )
    print(rds_warped1.meta)
    print(rds_warped1.data)
    rds_warped1.to_jp2("./test5_b08_wgs84_in_aoi.jp2")


def test_painting():
    fpath_input = os.path.join(
        BASE_PATH, "decloud_10m_51TXN_20200729_20200827_False.tif"
    )

    fpath_loc = os.path.join(
        BASE_PATH, "china_geojson_eng/shi4/qi2qi2ha1er3shi4.geojson"
    )
    info = json.loads(open(fpath_loc).read())
    aoi = Geom(info["features"][0]["geometry"])

    rds: RasterDataset = RasterDataset.from_file(fpath_input, bands=[0, 1, 2])
    print(rds.meta)
    print(rds.data.shape, rds.data.dtype)
    rds_painted = rds.paint(aoi, painted_values=0, inverse=True)
    rds_painted.to_geotiff("test8_aoi_painted.tif")


def test_cropping():
    pass


def test_create_dummy_image():
    fpath_input = os.path.join(
        BASE_PATH, "decloud_10m_51TXN_20200729_20200827_False.tif"
    )
    rst = RasterDataset.from_file(fpath_input)
    print(rst.meta)
    rst_dummy = RasterDataset.create_dummy_raster_like(
        rst.meta, dtype=np.int16, fill_value=-100
    )
    rst_dummy.to_geotiff("./test_dummy_raster.tif")


def test_from_geom():
    fpath_county = os.path.join(
        BASE_PATH, "china_geojson_eng/xian4/fu4yu4xian4.geojson"
    )
    geom_aoi = Geom.from_geojson_fpath(fpath_county)
    rst_mask = RasterDataset.from_geom(
        geom_aoi,
        dest_res=100,
        dest_srs=None,
        dtype=np.uint8,
        value_within_geom=10,
        value_outside_geom=0,
        nodata_val=0,
    )
    rst_mask.to_geotiff("./test_geom_mask.tif")


def test_stack_multiple_bands_together():
    dir_base = os.path.join(
        BASE_PATH, "S2B_MSIL2A_20190204T034949_N0211_R104_T48RVV_20190204T073411.SAFE"
    )
    fpath_input_blue = os.path.join(
        dir_base,
        "GRANULE/L2A_T48RVV_A009997_20190204T034943/IMG_DATA/R10m/T48RVV_20190204T034949_B02_10m.jp2",
    )
    fpath_input_green = os.path.join(
        dir_base,
        "GRANULE/L2A_T48RVV_A009997_20190204T034943/IMG_DATA/R10m/T48RVV_20190204T034949_B03_10m.jp2",
    )
    fpath_input_red = os.path.join(
        dir_base,
        "GRANULE/L2A_T48RVV_A009997_20190204T034943/IMG_DATA/R10m/T48RVV_20190204T034949_B04_10m.jp2",
    )
    fpath_input_nir = os.path.join(
        dir_base,
        "GRANULE/L2A_T48RVV_A009997_20190204T034943/IMG_DATA/R10m/T48RVV_20190204T034949_B08_10m.jp2",
    )
    fpath_loc = os.path.join(
        BASE_PATH, "china_geojson_eng/xian4/guang3han4shi4.geojson"
    )
    aoi = Geom.from_geojson_fpath(fpath_loc)

    with stack_ds(
        [
            fpath_input_blue,
            fpath_input_green,
            fpath_input_red,
            fpath_input_nir,
        ]
    ) as stacked_gdal_raster:
        rds_stacked = RasterDataset.from_gdal_raster(stacked_gdal_raster, aoi=aoi)
    rds_stacked.to_geotiff("test6_bgrn_stacked_s2.tif")


def test_mosaic_multiple_images_together():
    fpath_loc = os.path.join(BASE_PATH, "china_geojson_eng/xian4/fu4yu4xian4.geojson")
    aoi = Geom.from_geojson_fpath(fpath_loc)

    print("query to get the items")
    items = Sentinel2L2AScene.query_many_items(
        spatial=aoi,
        temporal=[datetime.datetime(2021, 7, 10), datetime.datetime(2021, 7, 11)],
    )
    items = list(items)

    print("mosaic ds")
    lst_fpaths = [
        interpret_protocol(item.info["assets"]["B08"]["href"]) for item in items
    ]
    print(lst_fpaths)
    with mosaic_ds(
        lst_fpaths,
        resample_alg="average",
        src_nodata=0,
        dest_res=100,
    ) as gdal_rst:
        rst = RasterDataset.from_gdal_raster(gdal_rst)
    rst.to_geotiff("test12_mosaiced_b08_100m.tif")
    del rst

    print("mosaic raster dataset")
    lst_rsts = [
        RasterDataset.from_file(interpret_protocol(item.info["assets"]["B04"]["href"]))
        for item in items
    ]
    mosaiced_rst = mosaic_raster_datasets(lst_rsts, resample_alg="nearest")
    mosaiced_rst.to_geotiff("test12_mosaiced_b04_10m.tif")


def main(option):
    if option == 1:
        test_open_single_band_image()
    elif option == 2:
        test_open_multiple_band_image()
    elif option == 3:
        test_open_multiple_band_image_real_case()
    elif option == 4:
        # test_iter_from_file(simple=False)
        test_iter_from_file(simple=True)
    elif option == 5:
        test_warp_to_wgs84()
    elif option == 6:
        test_stack_multiple_bands_together()
    elif option == 7:
        test_mosaic_multiple_images_together()
    elif option == 8:
        test_painting()
    elif option == 9:
        test_cropping()
    elif option == 10:
        test_create_dummy_image()
    elif option == 11:
        test_from_geom()
    elif option == 12:
        test_mosaic_multiple_images_together()
    elif option == "all":
        test_open_single_band_image()
        test_open_multiple_band_image()
        test_open_multiple_band_image_real_case()
        test_iter_from_file()
        test_warp_to_wgs84()


if __name__ == "__main__":
    import fire

    fire.Fire(main)
